import math
import torch
from torch import nn

from functools import partial

from visual_text_merge_test.weight_init import trunc_normal_


def drop_path(x, drop_prob: float = 0.0, training: bool = False):
    if drop_prob == 0 or not training:
        return x

    keep_prob = 1 - drop_prob
    shape = (x.shape[0],) + (1,) * (x.dim() - 1)
    rand_tensor = torch.rand(shape, dtype=x.dtype, device=x.device) + keep_prob
    mask = rand_tensor.floor()

    # 确保输出期望一致，这里需要除以keep_prob
    x = x.div(keep_prob) * mask
    return x


class DropPath(nn.Module):
    def __init__(self, drop_prob: float):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        return drop_path(x, self.drop_prob, self.training)


class MLP(nn.Module):
    def __init__(
        self,
        input_features: int,
        hidden_features: int = None,
        out_features: int = None,
        act_layer=nn.GELU,
        drop: float = 0.0,
    ):
        super().__init__()
        self.input_fearures = input_features
        self.hidden_features = hidden_features or input_features
        self.out_features = out_features if out_features else input_features

        self.fc1 = nn.Linear(self.input_fearures, self.hidden_features)
        self.act = act_layer()
        self.fc2 = nn.Linear(self.hidden_features, self.out_features)
        self.drop = nn.Dropout(drop)

    def forward(self, x):
        x = self.fc1(x)
        x = self.act(x)
        x = self.drop(x)
        x = self.fc2(x)
        x = self.drop(x)
        return x


class MultiheadSelfAttention(nn.Module):
    def __init__(
        self,
        dim,
        num_heads: int = 8,
        qkv_bias: bool = False,
        qk_scale=None,
        attn_drop=0.0,
        proj_drop=0.0,
    ):
        super().__init__()
        self.num_heads = num_heads
        head_dim = dim // num_heads
        self.scale = qk_scale or head_dim**-0.5

        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
        self.atten_drop = nn.Dropout(attn_drop)
        self.proj = nn.Linear(dim, dim)
        self.proj_drop = nn.Dropout(proj_drop)

    def forward(self, x):
        # (bacth_size,num_patches,dim)
        B, N, C = x.shape
        # (batch_size,num_patches,3,num_heads,dim_heads)->(3,batch_size,num_head,num_patches,dim_heads)
        qkv = (
            self.qkv(x)
            .reshape(B, N, 3, self.num_heads, C // self.num_heads)
            .permute(2, 0, 3, 1, 4)
        )
        # (batch_size,num_head,num_patches,dim_heads)
        q, k, v = qkv[0], qkv[1], qkv[2]

        # (batch_size,num_head,num_patches,dim_heads)@(batch_size,num_head,dim_heads,num_patches)
        # =(batch_size,num_head,num_patches,num_patches)
        atten = (self.scale * q) @ k.transpose(-2, -1)
        atten = atten.softmax(dim=-1)
        atten = self.atten_drop(atten)

        x = (atten @ v).transpose(1, 2).reshape(B, N, C)
        x = self.proj(x)
        x = self.proj_drop(x)
        return x, atten


class Block(nn.Module):
    def __init__(
        self,
        dim,
        num_heads,
        mlp_ratio=4.0,
        qkv_bias=False,
        qk_scale=None,
        drop=0.0,
        attn_drop=0.0,
        drop_path=0.0,
        act_layer=nn.GELU,
        norm_layer=nn.LayerNorm,
        init_values=0,
    ):
        super().__init__()
        self.norm1 = norm_layer(dim)
        self.attn = MultiheadSelfAttention(
            dim, num_heads, qkv_bias, qk_scale, attn_drop, drop
        )
        self.drop_path = DropPath(drop_path) if drop_path > 0 else nn.Identity()
        self.norm2 = norm_layer(dim)
        mlp_hidden_dim = dim * mlp_ratio
        self.mlp = MLP(dim, mlp_hidden_dim, act_layer=act_layer, drop=drop)
        if init_values > 0:
            self.gamma_1 = nn.Parameter(
                init_values * torch.ones(dim), requires_grad=True
            )
            self.gamma_2 = nn.Parameter(
                init_values * torch.ones(dim), requires_grad=True
            )
        else:
            self.gamma_1 = None
            self.gamma_2 = None

    def forward(self, x):
        x = self.norm1(x)
        x, atten = self.attn(x)
        # 残差连接
        if self.gamma_1 is None:
            x = x + self.drop_path(x)
            x = x + self.drop_path(self.mlp(self.norm2(x)))
        else:
            x = x + self.drop_path(self.gamma_1 * x)
            x = x + self.drop_path(self.gamma_2 * self.mlp(self.norm2(x)))
        return x


class PatchEmbed(nn.Module):
    def __init__(self, img_size=224, patch_size=16, in_chans=3, embed_dim=768):
        super().__init__()
        if isinstance(img_size, int):
            num_patches = (img_size // patch_size) * (img_size // patch_size)
        elif isinstance(img_size, tuple) or isinstance(img_size, list):
            if len(img_size) == 1:
                num_patches = (img_size[0] // patch_size) * (img_size[0] // patch_size)
            elif len(img_size) == 2:
                num_patches = (img_size[0] // patch_size) * (img_size[1] // patch_size)
            else:
                raise ValueError("img_size must be a tuple or list of length 1 or 2")
        else:
            raise ValueError("img_size must be a int, tuple or list")

        self.img_size = img_size
        self.patch_size = patch_size
        self.num_patches = num_patches
        self.proj = nn.Conv2d(
            in_channels=in_chans,
            out_channels=embed_dim,
            kernel_size=patch_size,
            stride=patch_size,
        )

    def forward(self, x):
        B, C, H, W = x.shape
        return self.proj(x)


class VisionTransformer(nn.Module):
    def __init__(
        self,
        img_size=[224],
        patch_size: int = 16,
        in_chans: int = 3,
        num_cls: int = 0,
        embed_dim: int = 768,
        depth: int = 12,
        num_heads: int = 12,
        mlp_ratio: float = 4.0,
        qkv_bias: bool = False,
        qk_scale=None,
        drop_rate: float = 0.0,
        attn_drop_rate: float = 0.0,
        drop_path_rate: float = 0.0,
        norm_layer=partial(nn.LayerNorm, eps=1e-6),
        return_all_tokens: bool = False,
        init_values=0,
        use_mean_pooling: bool = False,
    ):
        super().__init__()
        self.num_features = self.embed_dim = embed_dim
        self.return_all_tokens = return_all_tokens

        self.patch_embed = PatchEmbed(
            img_size=img_size,
            patch_size=patch_size,
            in_chans=in_chans,
            embed_dim=embed_dim,
        )
        num_patches = self.patch_embed.num_patches

        self.cls_token = nn.Parameter(torch.zeros(1, 1, embed_dim))
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 1, embed_dim))
        self.pos_drop = nn.Dropout(p=drop_rate)

        drop_path_rate_ls = [
            x.item() for x in torch.linspace(start=0, end=drop_path_rate, steps=depth)
        ]

        self.blocks = nn.ModuleList(
            [
                Block(
                    dim=embed_dim,
                    num_heads=num_heads,
                    mlp_ratio=mlp_ratio,
                    qkv_bias=qkv_bias,
                    qk_scale=qk_scale,
                    drop=drop_rate,
                    attn_drop=attn_drop_rate,
                    drop_path=drop_path_rate_ls[i],
                    norm_layer=norm_layer,
                    init_values=init_values,
                )
                for i in range(depth)
            ]
        )

        self.norm = nn.Identity() if use_mean_pooling else norm_layer(embed_dim)
        self.fc_norm = norm_layer(embed_dim, num_cls) if use_mean_pooling else None

        self.head = nn.Linear(embed_dim, num_cls) if num_cls > 0 else nn.Identity()

        trunc_normal_(self.pos_embed, std=0.02)
        trunc_normal_(self.cls_token, std=0.02)
        self.apply(self._init_weights)

    def _init_weights(self, m):
        if isinstance(m, nn.Linear):
            trunc_normal_(m.weight, std=0.02)
            if isinstance(m, nn.Linear) and m.bias is not None:
                nn.init.constant_(m.bias, 0)
        elif isinstance(m, nn.LayerNorm):
            nn.init.constant_(m.bias, 0)
            nn.init.constant_(m.weight, 1.0)

    def interpolate_pos_encoding(self, x, w, h):
        # x.shape: (B,feature_map_w*feature_map_h+1,embed_dim)
        # 1 为类别编码额外的维度
        # fw*fh
        npath = x.shape[1] - 1

        # self.pos_embed.shape: (1,num_patches+1,embed_dim)
        # num_patches
        N = self.pos_embed.shape[1] - 1

        # fw*fh==num_patches 且 刚好w==h,说明原始原始图像长宽刚好被patch_size整除，且为正方形,可以直接返还预定义的位置编码
        # 否则需要使用插值的方法获取位置编码
        if N == npath and w == h:
            return self.pos_embed

        # patch_pos_embed.shape:(1,1,embed_dim)
        class_pos_embed = self.pos_embed[:, 0]
        # patch_pos_embed.shape:(1,num_patches,embed_dim)
        patch_pos_embed = self.pos_embed[:, 1:]

        # embed_dim
        dim = x.shape[-1]

        # feature map size
        h0 = h // self.patch_embed.patch_size
        w0 = w // self.patch_embed.patch_size
        w0, h0 = w0 + 0.1, h0 + 0.1

        perdefine_feature_w = perdefine_feature_h = int(math.sqrt(N))

        # patch_pos_embed.shape:
        # (1,num_patches,embed_dim)
        # -reshape->
        # (1,perdefine_feature_w,perdefine_feature_h,embed_dim)
        # -permute(0,3,1,2)->
        # (1,embed_dim,perdefine_feature_w,perdefine_feature_h)
        # -interpolate->
        # (1,embed_dim,w0,h0)
        patch_pos_embed = nn.functional.interpolate(
            input=patch_pos_embed.reshape(
                1, perdefine_feature_w, perdefine_feature_h, dim
            ).permute(0, 3, 1, 2),
            scale_factor=(w0 / perdefine_feature_w, h0 / perdefine_feature_h),
            mode="bicubic",
        )

        # 确保线性插值后输出结果尺寸与要求尺寸一致
        assert (
            int(w0) == patch_pos_embed.shape[-2]
            and int(h0) == patch_pos_embed.shape[-1]
        )

        # patch_pos_embed.shape: (1,embed_dim,w0,h0)->(1,w0,h0,embed_dim)->(1,w0*h0,embed_dim)
        patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
        return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)

    def prepare_tokens(self, x, mask=None):
        B, nc, w, h = x.shape
        # feture_map: 指代图像在经过patch_size大小卷积处理后，输出的张量
        # x.shape = (B,embed_dim,feature_map_w,feature_map_h)
        x = self.patch_embed(x)

        if mask is not None:
            x = self.mask_model(x, mask)

        # (B,embed_dim,fw,fh) -> (B,embed_dim,fw*fh) -> (B,fw*fh,embed_dim)
        x = x.flatten(2).transpose(1, 2)

        # cls_tokens.shape=(B,1,embed_dim)
        cls_tokens = self.cls_token.expand(B, -1, -1)

        # x.shape: (B,fw*fh,embed_dim)->(B,fw*fh+1,embed_dim)
        x = torch.cat((cls_tokens, x), dim=1)

        # x.shape: (B,fw*fh,embed_dim)->(B,fw*fh+1,embed_dim)
        x = x + self.interpolate_pos_encoding(x, w, h)

        return self.pos_drop(x)

    def forward(self, x, return_all_tokens=None):
        # x.shape:(B,C,H,W)->(B,fw*fh+1,embed_dim)
        x = self.prepare_tokens(x)

        # x.shape:(B,fw*fh+1,embed_dim)->(B,fw*fh+1,embed_dim)
        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x)
        if self.fc_norm is not None:
            x[:, 0] = self.fc_norm(x[:, 1:, :].mean(1))

        return_all_tokens = (
            return_all_tokens
            if return_all_tokens is not None
            else self.return_all_tokens
        )
        if return_all_tokens:
            return x

        return x[:, 0]


def vit_tiny(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size,
        embed_dim=192,
        depth=12,
        num_heads=3,
        mlp_ratio=4,
        qkv_bias=True,
        **kwargs
    )
    return model


def vit_small(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size,
        embed_dim=384,
        depth=12,
        num_heads=6,
        mlp_ratio=4,
        qkv_bias=True,
        **kwargs
    )
    return model


def vit_base(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size,
        embed_dim=768,
        depth=12,
        num_heads=12,
        mlp_ratio=4,
        qkv_bias=True,
        **kwargs
    )
    return model


def vit_large(patch_size=16, **kwargs):
    model = VisionTransformer(
        patch_size=patch_size,
        embed_dim=1024,
        depth=24,
        num_heads=16,
        mlp_ratio=4,
        qkv_bias=True,
        **kwargs
    )
    return model
